import os
import time
import numpy as np
from sympy import primerange, nextprime
from concurrent.futures import ProcessPoolExecutor, as_completed
import multiprocessing


def get_last_prime_from_file(file_name, dtype):
    """从二进制文件中读取最后一个素数"""
    try:
        with open(file_name, 'rb') as file:
            file.seek(-np.dtype(dtype).itemsize, os.SEEK_END)
            last_prime_bytes = file.read(np.dtype(dtype).itemsize)
            return int.from_bytes(last_prime_bytes, byteorder='little', signed=True)
    except (FileNotFoundError, OSError):
        return 1
    except Exception as e:
        print(f"读取文件时发生未知错误: {e}")
        return 1


def generate_primes(start_prime, count):
    """生成从start_prime开始的count个素数"""
    return list(primerange(start_prime, start_prime + count))


def append_primes_to_binary(file_name, total_count, data_size=4, max_workers=None):
    """将生成的素数追加到二进制文件中"""
    if data_size not in (4, 8):
        raise ValueError("数据大小必须是4（int32）或8（int64）。")

    dtype = np.int32 if data_size == 4 else np.int64
    start_prime = get_last_prime_from_file(file_name, dtype)
    start_prime = nextprime(start_prime)

    chunk_size = max(total_count // 100, 1)  # 至少为1
    extra_primes = total_count % 100
    num_workers = max_workers if max_workers is not None else os.cpu_count()

    with ProcessPoolExecutor(max_workers=num_workers) as executor:
        futures = {}
        for i in range(100):
            future = executor.submit(generate_primes, start_prime + i * chunk_size,
                                     chunk_size + (1 if i < extra_primes else 0))
            futures[future] = i + 1

        primes_list = []
        for future in as_completed(futures):
            try:
                primes_list.extend(future.result())
                if futures[future] % 10 == 0:
                    print(f"进度：{futures[future]}% 的块已完成")
            except Exception as exc:
                print(f"素数生成过程中发生错误：{exc}")

    primes_array = np.array(primes_list, dtype=dtype)

    try:
        with open(file_name, 'ab') as f:
            # 使用更大的缓冲区可能提高写入速度
            f.write(primes_array.tobytes())
    except Exception as e:
        print(f"写入文件时发生错误：{e}")


def print_primes_from_file(file_name, count, data_size):
    """从二进制文件中读取并打印素数"""
    if data_size not in (4, 8):
        raise ValueError("数据大小必须是4（int32）或8（int64）。")

    dtype = np.int32 if data_size == 4 else np.int64

    try:
        with open(file_name, 'rb') as f:
            primes = np.fromfile(f, dtype=dtype, count=count) if count != -1 else np.fromfile(f, dtype=dtype)
            for prime in primes:
                print(prime)
    except FileNotFoundError:
        print(f"文件 {file_name} 未找到。")
    except Exception as e:
        print(f"读取文件时发生错误：{e}")


if __name__ == '__main__':
    multiprocessing.freeze_support()

    n = int(1e7)
    for _ in range(40):
        start_time = time.time()
        append_primes_to_binary('primes.bin', n, 8)
        print(f"耗时：{time.time() - start_time}秒")

        print(f"最后一个素数：{get_last_prime_from_file('primes.bin', np.int64)}")
